#include "buffer.h"
#include "common.h"
#include <stdint.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

int main() {
  aclInit(NULL);
  int ret = buffer_init(DATA_SIZE);
  if (ret != 0) {
    fprintf(stderr, "buffer_init failed\n");
    return -1;
  }
  aclrtStream d2hStream, h2dStream;
  ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
  ACL_CHECK(aclrtCreateStream(&d2hStream));
  ACL_CHECK(aclrtSetDevice(TO_DEVICE));
  ACL_CHECK(aclrtCreateStream(&h2dStream));

  aclrtEvent d2hStartEvent, d2hEndEvent, h2dStartEvent, h2dEndEvent;
  ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
  ACL_CHECK(aclrtCreateEventWithFlag(&d2hStartEvent, ACL_EVENT_TIME_LINE));
  ACL_CHECK(aclrtCreateEventWithFlag(&d2hEndEvent,
                                     ACL_EVENT_TIME_LINE | ACL_EVENT_SYNC));
  ACL_CHECK(aclrtSetDevice(TO_DEVICE));
  ACL_CHECK(aclrtCreateEventWithFlag(&h2dStartEvent, ACL_EVENT_TIME_LINE));
  ACL_CHECK(aclrtCreateEventWithFlag(&h2dEndEvent, ACL_EVENT_TIME_LINE));

  // Warm-up iterations
  for (int i = 0; i < WARMUP_ITERATIONS; i++) {
    ACL_CHECK(aclrtSetDevice(TO_DEVICE));
    ACL_CHECK(aclrtMemset(dstbuf_d, DATA_SIZE, 0, DATA_SIZE));
    memset(relaybuf_h, 0, DATA_SIZE);
    ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
    ACL_CHECK(aclrtMemcpyAsync(relaybuf_h, DATA_SIZE, srcbuf_d, DATA_SIZE,
                               ACL_MEMCPY_DEVICE_TO_HOST, d2hStream));
    ACL_CHECK(aclrtRecordEvent(d2hEndEvent, d2hStream));
    ACL_CHECK(aclrtSynchronizeStream(d2hStream)); // WHY
    ACL_CHECK(aclrtSetDevice(TO_DEVICE));
    ACL_CHECK(aclrtStreamWaitEvent(h2dStream, d2hEndEvent));
    ACL_CHECK(aclrtMemcpyAsync(dstbuf_d, DATA_SIZE, relaybuf_h, DATA_SIZE,
                               ACL_MEMCPY_HOST_TO_DEVICE, h2dStream));
    ACL_CHECK(aclrtSynchronizeStream(h2dStream));
    if (buffer_verify() != 0) {
      fprintf(stderr, "data verify failed\n");
      return -1;
    }
  }

  // Measured iterations
  float totalD2hElapsedTime = 0.0f, totalH2dElapsedTime = 0.0f;
  double totalWallTime = 0.0;
  for (int i = 0; i < MEASURED_ITERATIONS; i++) {
    float d2hElapsedTime = 0.0f, h2dElapsedTime = 0.0f;
    struct timespec startWall, endWall;
    ACL_CHECK(aclrtSetDevice(TO_DEVICE));
    ACL_CHECK(aclrtMemset(dstbuf_d, DATA_SIZE, 0, DATA_SIZE));
    memset(relaybuf_h, 0, DATA_SIZE);
    ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
    clock_gettime(CLOCK_MONOTONIC, &startWall);
    ACL_CHECK(aclrtRecordEvent(d2hStartEvent, d2hStream));
    ACL_CHECK(aclrtMemcpyAsync(relaybuf_h, DATA_SIZE, srcbuf_d, DATA_SIZE,
                               ACL_MEMCPY_DEVICE_TO_HOST, d2hStream));
    ACL_CHECK(aclrtRecordEvent(d2hEndEvent, d2hStream));
    ACL_CHECK(aclrtSynchronizeStream(d2hStream)); // WHY
    ACL_CHECK(aclrtSetDevice(TO_DEVICE));
    ACL_CHECK(aclrtStreamWaitEvent(h2dStream, d2hEndEvent));
    ACL_CHECK(aclrtRecordEvent(h2dStartEvent, h2dStream));
    ACL_CHECK(aclrtMemcpyAsync(dstbuf_d, DATA_SIZE, relaybuf_h, DATA_SIZE,
                               ACL_MEMCPY_HOST_TO_DEVICE, h2dStream));
    ACL_CHECK(aclrtRecordEvent(h2dEndEvent, h2dStream));

    ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
    ACL_CHECK(aclrtSynchronizeStream(d2hStream));
    ACL_CHECK(aclrtSetDevice(TO_DEVICE));
    ACL_CHECK(aclrtSynchronizeStream(h2dStream));
    clock_gettime(CLOCK_MONOTONIC, &endWall);

    ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
    ACL_CHECK(
        aclrtEventElapsedTime(&d2hElapsedTime, d2hStartEvent, d2hEndEvent));
    ACL_CHECK(aclrtSetDevice(TO_DEVICE));
    ACL_CHECK(
        aclrtEventElapsedTime(&h2dElapsedTime, h2dStartEvent, h2dEndEvent));

    double wallTime = (endWall.tv_sec - startWall.tv_sec) * 1000.0 +
                      (endWall.tv_nsec - startWall.tv_nsec) / 1e6;
    totalWallTime += wallTime;
    if (buffer_verify() != 0) {
      fprintf(stderr, "data verify failed\n");
      return -1;
    }
    totalD2hElapsedTime += d2hElapsedTime;
    totalH2dElapsedTime += h2dElapsedTime;
  }
  totalD2hElapsedTime /= MEASURED_ITERATIONS;
  totalH2dElapsedTime /= MEASURED_ITERATIONS;
  totalWallTime /= MEASURED_ITERATIONS;

  printf("Average D2H Time: %.3f ms, Bandwidth: %.3f GB/s\n",
         totalD2hElapsedTime,
         ((float)DATA_SIZE) / (totalD2hElapsedTime * 1e6f));
  printf("Average H2D Time: %.3f ms, Bandwidth: %.3f GB/s\n",
         totalH2dElapsedTime,
         ((float)DATA_SIZE) / (totalH2dElapsedTime * 1e6f));

  float totalElapsedTime = totalD2hElapsedTime + totalH2dElapsedTime;
  printf("Combined Time: %.3f ms, Bandwidth: %.3f GB/s\n", totalElapsedTime,
         ((float)DATA_SIZE) / (totalElapsedTime * 1e6f));
  printf("Wall Clock Time: %.3f ms, Bandwidth: %.3f GB/s\n", totalWallTime,
         ((float)DATA_SIZE) / (totalWallTime * 1e6));
  ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
  ACL_CHECK(aclrtDestroyEvent(d2hStartEvent));
  ACL_CHECK(aclrtDestroyEvent(d2hEndEvent));
  ACL_CHECK(aclrtDestroyStream(d2hStream));
  ACL_CHECK(aclrtSetDevice(TO_DEVICE));
  ACL_CHECK(aclrtDestroyEvent(h2dStartEvent));
  ACL_CHECK(aclrtDestroyEvent(h2dEndEvent));
  ACL_CHECK(aclrtDestroyStream(h2dStream));
  ret = buffer_destroy();
  if (ret != 0) {
    fprintf(stderr, "buffer_destroy failed\n");
    return -1;
  }
  aclFinalize();
  return 0;
}